#include "compi/compi.h"

#include <iostream>
#include <string>

using namespace compi;

// 使用合并后的 Channel 泛型 API（send/recv<T>）进行点对点往返演示。
static void
demo_channel_typed(Context & ctx)
{
    int rank = ctx.rank();
    int peer = (rank == 0 ? 1 : 0);

    auto & mgr = ChannelManager::for_comm(ctx.comm());
    auto ch_res = mgr.create_channel(ChannelId("typed_merged"), /*use_dedicated_comm=*/true);
    if (ch_res.is_error())
    {
        if (rank == 0)
            std::cerr << "create channel error: " << ch_res.error_message() << std::endl;
        return;
    }
    auto ch = ch_res.value();
    (void)ch->activate();

    // 使用命名空间进行 ANY_TAG 过滤
    auto & alloc = ctx.tag_allocator();
    TagNamespace ns = alloc.create_namespace("typed_demo_ns");

    if (rank == 0)
    {
        // 1) 显式 tag 的整数往返
        int send_val = 1234;
        auto s1 = ch->send<int>(peer, send_val, /*tag=*/101);
        if (s1.is_error())
            std::cerr << "[channel-typed] send int error: " << s1.error_message() << std::endl;

        int recv_val = 0;
        auto r1 = ch->recv<int>(peer, recv_val, /*tag=*/102);
        if (r1.is_error())
            std::cerr << "[channel-typed] recv int error: " << r1.error_message() << std::endl;
        else
            std::cout << "[channel-typed] rank0 recv int: " << recv_val << std::endl;

        // 2) 命名空间（ANY_TAG 过滤）字符串往返
        std::string msg = "hello typed";
        auto s2 = ch->send<std::string>(peer, msg, ns);
        if (s2.is_error())
            std::cerr << "[channel-typed] send str error: " << s2.error_message() << std::endl;

        std::string reply;
        auto r2 = ch->recv<std::string>(peer, reply, ns);
        if (r2.is_error())
            std::cerr << "[channel-typed] recv str error: " << r2.error_message() << std::endl;
        else
            std::cout << "[channel-typed] rank0 recv str: " << reply << std::endl;

        // 3) 发送一个显式标签的最终消息，供对端通过 iprobe 探测后按确切标签接收
        TagLease lease = ns.lease();
        std::string final_msg = "final-by-lease";
        auto s3 = ch->send<std::string>(peer, final_msg, /*tag=*/lease.value());
        if (s3.is_error())
            std::cerr << "[channel-typed-lease] send str error: " << s3.error_message()
                      << std::endl;
        else
            std::cout << "[channel-typed-lease] rank0 sent with tag " << lease.value() << ": "
                      << final_msg << std::endl;
        lease.retire();
    }
    else
    {
        // 1) 显式 tag 的整数往返（接收后发送）
        int got = 0;
        auto r1 = ch->recv<int>(peer, got, /*tag=*/101);
        if (r1.is_error())
            std::cerr << "[channel-typed] recv int error: " << r1.error_message() << std::endl;
        else
            std::cout << "[channel-typed] rank1 recv int: " << got << std::endl;

        auto s1 = ch->send<int>(peer, got + 1, /*tag=*/102);
        if (s1.is_error())
            std::cerr << "[channel-typed] send int error: " << s1.error_message() << std::endl;

        // 2) 命名空间（ANY_TAG 过滤）字符串往返
        std::string got_str;
        auto r2 = ch->recv<std::string>(peer, got_str, ns);
        if (r2.is_error())
            std::cerr << "[channel-typed] recv str error: " << r2.error_message() << std::endl;
        else
            std::cout << "[channel-typed] rank1 recv str: " << got_str << std::endl;

        auto s2 = ch->send<std::string>(peer, got_str + " | ack", ns);
        if (s2.is_error())
            std::cerr << "[channel-typed] send str error: " << s2.error_message() << std::endl;

        // 3) 命名空间 iprobe 获取确切标签后再显式接收
        MessageTag probed_tag(0, ns.id());
        bool ready = false;
        while (!ready)
        {
            // 重要：使用 Channel 的 communicator 进行探测，避免与专用 communicator 不一致导致卡死
            ready = ns.iprobe(ch->comm(), peer, probed_tag);
        }

        std::string lgot;
        auto r3 = ch->recv<std::string>(peer, lgot, /*tag=*/probed_tag.value);
        if (r3.is_error())
            std::cerr << "[channel-typed-lease] recv str error: " << r3.error_message()
                      << std::endl;
        else
            std::cout << "[channel-typed-lease] rank1 recv with tag " << probed_tag.value << ": "
                      << lgot << std::endl;
    }

    (void)ch->close();
    (void)mgr.destroy_channel(ChannelId("typed_merged"));
}

int
main()
{
    Environment env;
    env.set_failure_policy(FailurePolicy::ThrowException);
    Context & ctx = Context::for_comm(MPI_COMM_WORLD, env);

    if (ctx.size() < 2)
    {
        if (ctx.rank() == 0)
            std::cerr << "This example requires at least 2 MPI processes." << std::endl;
        return 0;
    }

    demo_channel_typed(ctx);
    return 0;
}